package com.yh.ldzl.config.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 质检WebSocket处理器
 * 负责处理与质检相关的WebSocket连接和消息
 */
@Component
public class QualityCheckWebSocketHandler extends TextWebSocketHandler {

    // 存储会话ID到WebSocketSession的映射
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();
    // 存储用户ID到会话ID的映射
    private final Map<String, String> userToSessionId = new ConcurrentHashMap<>();
    
    private final ObjectMapper objectMapper = new ObjectMapper();
    
    /**
     * 连接建立后调用
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        System.out.println("WebSocket连接已建立 - URI: " + session.getUri() + ", ID: " + session.getId());
        // 获取用户ID
        String userId = getUserId(session);
        String sessionId = session.getId();
        
        // 存储会话
        sessions.put(sessionId, session);
        if (userId != null) {
            userToSessionId.put(userId, sessionId);
        }
        
        // 发送连接成功消息
        WebSocketMessage<String> message = new WebSocketMessage<>(
            "connected", 
            "WebSocket连接成功",
            "欢迎使用质检通知系统"
        );
        
        sendMessage(session, message);
        
        System.out.println("质检WebSocket连接已建立 - 会话ID: " + sessionId + ", 用户ID: " + userId);
    }

    /**
     * 处理文本消息
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        
        try {
            // 处理接收到的消息
            System.out.println("收到WebSocket消息: " + payload);
            
            // 处理心跳包
            if (payload.contains("ping")) {
                WebSocketMessage<String> response = new WebSocketMessage<>(
                    "ping_response",
                    "pong",
                    "心跳响应"
                );
                sendMessage(session, response);
                return;
            }
            
            // 可以在这里处理其他类型的消息
            // ...
            
        } catch (Exception e) {
            System.err.println("处理WebSocket消息出错: " + e.getMessage());
            // 发送错误响应
            WebSocketMessage<String> errorMessage = new WebSocketMessage<>(
                "error",
                "处理消息时发生错误",
                e.getMessage()
            );
            sendMessage(session, errorMessage);
        }
    }

    /**
     * 处理传输错误
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        System.out.println("WebSocket传输错误 - URI: " + session.getUri() + ", ID: " + session.getId() + ", 错误: " + exception.getMessage());
        System.err.println("WebSocket传输错误: " + exception.getMessage());
        if (session.isOpen()) {
            session.close();
        }
        removeSession(session);
    }

    /**
     * 连接关闭后调用
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        System.out.println("WebSocket连接已关闭 - URI: " + session.getUri() + ", ID: " + session.getId() + ", 状态: " + closeStatus);
        removeSession(session);
        System.out.println("WebSocket连接已关闭 - 会话ID: " + session.getId() + ", 状态: " + closeStatus);
    }
    
    /**
     * 广播消息给所有连接的客户端
     */
    public <T> void broadcastMessage(WebSocketMessage<T> message) {
        String jsonMessage;
        try {
            jsonMessage = objectMapper.writeValueAsString(message);
            
            for (WebSocketSession session : sessions.values()) {
                if (session.isOpen()) {
                    try {
                        session.sendMessage(new TextMessage(jsonMessage));
                    } catch (IOException e) {
                        System.err.println("发送消息失败: " + e.getMessage());
                    }
                }
            }
        } catch (Exception e) {
            System.err.println("广播消息序列化失败: " + e.getMessage());
        }
    }
    
    /**
     * 发送消息给特定用户
     */
    public <T> boolean sendMessageToUser(String userId, WebSocketMessage<T> message) {
        String sessionId = userToSessionId.get(userId);
        if (sessionId == null) {
            return false;
        }
        
        WebSocketSession session = sessions.get(sessionId);
        if (session == null || !session.isOpen()) {
            return false;
        }
        
        try {
            sendMessage(session, message);
            return true;
        } catch (Exception e) {
            System.err.println("发送消息给用户失败: " + e.getMessage());
            return false;
        }
    }
    
    /**
     * 向WebSocketSession发送消息
     */
    private <T> void sendMessage(WebSocketSession session, WebSocketMessage<T> message) throws IOException {
        if (session.isOpen()) {
            String jsonMessage = objectMapper.writeValueAsString(message);
            session.sendMessage(new TextMessage(jsonMessage));
        }
    }
    
    /**
     * 获取当前连接数
     */
    public int getConnectionCount() {
        return sessions.size();
    }
    
    /**
     * 从会话列表中移除会话
     */
    private void removeSession(WebSocketSession session) {
        String sessionId = session.getId();
        String userId = getUserId(session);
        
        sessions.remove(sessionId);
        if (userId != null) {
            userToSessionId.remove(userId);
        }
    }
    
    /**
     * 从会话中获取用户ID
     */
    private String getUserId(WebSocketSession session) {
        Object userId = session.getAttributes().get("userId");
        return userId != null ? userId.toString() : null;
    }
}
